/*
 * anls_concensus.cc
 *
 * Copyright 2013 Long Tran (ltran3@gatech.edu)
 *
 * Master/Slave ALS, solving saddle point of the Lagrangian
 *   L = sum |A_i - W_iH_i^T|^2 + <Lambda_i, W_i-U_i> + <Pi_i, H_i-Z>
 *        + 0.5*rho*(|W_i-U_i|^2+|H_i-Z|^2)
 *   subject to U_i >= 0, Z >= 0
 * 
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 * 
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 * 
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
 * MA 02110-1301, USA.
 * 
 */

#include <armadillo>
#include <boost/archive/text_oarchive.hpp>
#include <boost/archive/text_iarchive.hpp>
#include <boost/mpi.hpp>

#include "helper/helper.h"
#include "algorithm/StopWatch.h"
#include "data/FileAccess.h"
extern boost::mpi::communicator* p_world;
extern po::variables_map vm;
#include "mpi_helper.h"
#include "arma_helper.h"

#define MSG_TAG_A 0

class ConcensusANLS
{
public:
  // input/output variables
  arma::mat A, W, H;
  // iteration information
  /*
  struct Iteration
  {
    int iter;
    double residual;
    TimeDuration duration, compute, communication;
    Iteration(int iter_, double residual_) : iter(iter_), residual(residual_) {}
  };
  */
  std::vector<Iteration> iterInfo;
protected:
  // algorithm variables
  int m, n, k;
  double rho, alpha;
  double l1, l2;
  int maxiter, interval;
  bool normalizing;
  arma::mat U, Z, Lambda, Pi, U_old, Z_old;
  double rho_incr, rho_decr;

  StopWatch watch, clock;
  TimeDuration runtime[2], total_time[2];
public:
  ConcensusANLS()
  {
    PO_REQUIRE(vm, "iter");

    PO_VARIABLE_VALUE(rho, vm, "rho", double); // ADMM parameter
    PO_VARIABLE_VALUE(rho_incr, vm, "rho-incr", double); // ADMM parameter
    PO_VARIABLE_VALUE(rho_decr, vm, "rho-decr", double); // ADMM parameter
    PO_VARIABLE_VALUE(alpha, vm, "alpha", double); // consensus penalty
    l1 = (vm.count("l1")) ? vm["l1"].as<double>() : 0; // l1 regularization
    l2 = (vm.count("l2")) ? vm["l2"].as<double>() : 0; // l2 regularization
    PO_VARIABLE_VALUE(maxiter, vm, "iter", int);
    PO_VARIABLE_VALUE(interval, vm, "interval", int);
    normalizing = vm.count("normalize");
  }

  void LoadData() 
  {
    PO_REQUIRE(vm, "data");
    PO_REQUIRE(vm, "k");
    PO_VARIABLE(data, vm, "data", std::string);
    PO_VARIABLE_VALUE(k, vm, "k", int);
    
    using namespace arma;
    std::stringstream ss;
    ss << data << "_" << MPI_RANK;
    A.load(ss.str(), arma_binary);
    m = A.n_rows;
    n = A.n_cols;
  }

  void UpdateWHU() 
  {
    using namespace arma;
    W = trans(solve(trans(H)*H+rho*eye(k,k), trans(A*H-Lambda+rho*U)));
    H = trans(solve(trans(W)*W+alpha*eye(k,k), trans(trans(A)*W-Pi+alpha*Z)));

    U_old = U;
    U = W + Lambda/rho; 
    zero_threshold(U);
  }

  void UpdateZ()
  {
    using namespace arma;
    //mat sum_H, sum_Pi;
    mat HpPi = H*alpha + Pi;
    mat sum;
    //mpi::reduce(*p_world, H, sum_H, std::plus<mat>(), 0);
    //mpi::reduce(*p_world, Pi, sum_Pi, std::plus<mat>(), 0);
    mpi::reduce(*p_world, HpPi, sum, std::plus<mat>(), 0);
    if (MPI_RANK == 0) {
      Z_old = Z;
      //MSG("sum HP reduced");
      //Z = sum_H / MPI_SIZE + sum_Pi / (MPI_SIZE*rho);
      sum -= l1*ones<mat>(n,k);
      double gamma = alpha*MPI_SIZE+l2;
      Z = sum / gamma;
      if (normalizing)
        normalize_column_equal_1(Z);
      else
        zero_threshold(Z);
    }
    mpi::broadcast(*p_world, Z, 0);
  }
  
  void UpdateLambdaPi()
  {
    //double rate = /*alpha / sqrt(iter+1)**/rho;
    Lambda += rho*(W-U);
    Pi += alpha*(H-Z);
  }
  
  void UpdateRho()
  {
    // primal feasibility
    double w_norm = norm(W-U, "fro"), h_norm = norm(H-Z, "fro"), sum_w, sum_h;
    w_norm *= w_norm;
    h_norm *= h_norm;
    mpi::reduce(*p_world, w_norm, sum_w, std::plus<double>(), 0);
    mpi::reduce(*p_world, h_norm, sum_h, std::plus<double>(), 0);
    
    double u_norm = norm( rho*(U_old - U), "fro" ), sum_u; 
    u_norm *= u_norm;
    mpi::reduce(*p_world, u_norm, sum_u, std::plus<double>(), 0);

    if (MPI_RANK == 0) {
      double v_norm = norm(alpha*(Z-Z_old), "fro"), sum_v = v_norm*v_norm*MPI_SIZE;
      
      if (sum_w > 100*sum_u) rho *= rho_incr;
      else if (sum_u > 100*sum_w) rho *= rho_decr;

      if (sum_h > 100*sum_v) alpha *= rho_incr;
      else if (sum_v > 100*sum_h) alpha *= rho_decr;
      
      if (rho > 1e3) rho = 1e3;
      if (rho < 1e-2) rho = 1e-2;  
      if (alpha > 1e3) alpha = 1e3;
      if (alpha < 1e-2) alpha = 1e-2;
    }
    mpi::broadcast(*p_world, rho, 0);
    mpi::broadcast(*p_world, alpha, 0);
  }

  void RecordProgress(int iter)
  {
    double res = norm(A-U*trans(Z), "fro"), sum_res;
    res *= res;
    mpi::reduce(*p_world, res, sum_res, std::plus<double>(), 0);
    if (MPI_RANK == 0) {
      Iteration info(iter, sqrt(sum_res), runtime[0]+runtime[1], runtime[0], runtime[1]);
      iterInfo.push_back(info);
      if (iter % interval == 0)
	MSG("it=" << iter << " |res|=" << sqrt(sum_res)  << " |H|=" << norm(Z, "fro")
	    << " d=" << SECONDS(info.duration)
	    << " alpha=" << alpha << " rho=" << rho);
    }
  }

  void SolveANLS()
  {
    using namespace arma;
    W = randu<mat>(m, k); 
    H = randu<mat>(n, k); 

    U = randu<mat>(m,k);
    Z = randu<mat>(n,k);
    mpi::broadcast(*p_world, Z, 0);
    Lambda = zeros<mat>(m,k);
    Pi = zeros<mat>(n,k);

    watch.Reset();
    clock.Reset();
    iterInfo.clear();
    for (int iter = 0; iter < maxiter; iter++) {
      RecordProgress(iter);
      // update W_i, H_i, U
      clock.GetDuration();
      UpdateWHU();
      runtime[0] += clock.GetDuration();
    
      // update Z
      clock.GetDuration();
      UpdateZ();
      runtime[1] += clock.GetDuration();

      // update Lambda_i, Pi_i
      clock.GetDuration();
      UpdateLambdaPi();
      runtime[0] += clock.GetDuration();
    
      UpdateRho();
    }
    RecordProgress(maxiter);
    if (MPI_RANK == 0) {
      MSG(watch.GetTotalDuration());
    }
    mpi::reduce(*p_world, runtime[0], total_time[0], std::plus<TimeDuration>(), 0);
    mpi::reduce(*p_world, runtime[1], total_time[1], std::plus<TimeDuration>(), 0);
    if (MPI_RANK == 0)
      for (int i = 0; i < 2; i++)
        MSG("clock " << i << " " << total_time[i]);
    W = U;
    H = Z;    
  }
};

void run_master_slave_anls()
{
  if (MPI_RANK == 0) MSG("Master/Slave ANLS 123");
  ConcensusANLS a;
  a.LoadData();
  a.SolveANLS();
  
  if (MPI_RANK == 0 && vm.count("output")) {
    PO_VARIABLE(output, vm, "output", std::string);
    a.H.save(output, arma::raw_ascii);
  }

  if (MPI_RANK == 0 && vm.count("iterinfo")) {
    PO_VARIABLE(infoFile, vm, "iterinfo", std::string);
    FileAccessWrite(out, infoFile, NULL);
    const std::vector<Iteration>& info =  a.iterInfo;
    for (unsigned int i = 0; i < info.size(); i++) {
      out << info[i].iter << "\t" << SECONDS(info[i].duration) << "\t" 
	  << SECONDS(info[i].compute) << "\t" 
	  << SECONDS(info[i].communication) << "\t"
	  << info[i].residual << std::endl;
    }
  }
}

/*
void solve_master_slave_anls(const arma::mat& A, arma::mat& W, arma::mat& H)
{
  StopWatch watch;
  PO_REQUIRE(vm, "iter");

  PO_VARIABLE(rho, vm, "rho", double); // ADMM parameter
  PO_VARIABLE(alpha, vm, "rate", double); // dual ascent step
  PO_VARIABLE(maxiter, vm, "iter", int);

  int m = A.n_rows, n = A.n_cols, k = W.n_cols; // local number of rows (not the total)
  using namespace arma;
  mat U = zeros<mat>(m,k), Z = zeros<mat>(n,k); // Z is computed in master then broadcasted to all
  mat Lambda = zeros<mat>(m,k), Pi = zeros<mat>(n,k);
  
  //MSG_MASTER(1, "A=\n" << A << " U=\n" << U);
  
  StopWatch clocks[2];
  TimeDuration runtime[2];
  for (int iter = 0; iter < maxiter; iter++) {
    // update W_i, H_i
    clocks[0].GetDuration();
    mat AA = trans(H)*H+rho*eye(k,k);
    mat bb = A*H-Lambda+rho*U;
    W = trans(solve(AA, trans(bb)));
    
    AA = trans(W)*W+rho*eye(k,k);
    bb = trans(A)*W-Pi+rho*Z;
    H = trans(solve(AA, trans(bb)));
    
    // update U_i
    U = W + Lambda/rho; 
    zero_threshold(U);
    runtime[0] += clocks[0].GetDuration();
    
    // update Z
    // reduce sum of H_i, Pi_i
    clocks[1].GetDuration();
    mat sum_H, sum_Pi;
    mpi::reduce(*p_world, H, sum_H, std::plus<mat>(), 0);
    mpi::reduce(*p_world, Pi, sum_Pi, std::plus<mat>(), 0);
    if (MPI_RANK == 0) {
      Z = sum_H / MPI_SIZE + sum_Pi / (MPI_SIZE*rho);
      zero_threshold(Z);
    }
    mpi::broadcast(*p_world, Z, 0);
    runtime[1] += clocks[1].GetDuration();

    // update Lambda_i, Pi_i
    clocks[0].GetDuration();
    double rate = alpha / sqrt(iter+1)*rho;
    Lambda += rate*(W-U);
    Pi += rate*(H-Z);
    runtime[0] += clocks[0].GetDuration();
    
    if (iter % 100 == 0 || iter == maxiter-1) {
      double res = norm(A-U*trans(Z), "fro"), sum_res;
      res *= res;
      mpi::reduce(*p_world, res, sum_res, std::plus<double>(), 0);
      if (MPI_RANK == 0)
        MSG("it=" << iter << " |res|=" << sqrt(sum_res)  << " norm(H,fro)=" << norm(Z, "fro")
            << " d=" << watch.GetDuration());
    }
  }
  if (MPI_RANK == 0) {
    MSG(watch.GetTotalDuration());
  }
  TimeDuration total[2];
  mpi::reduce(*p_world, runtime[0], total[0], std::plus<TimeDuration>(), 0);
  mpi::reduce(*p_world, runtime[1], total[1], std::plus<TimeDuration>(), 0);
  if (MPI_RANK == 0)
    for (int i = 0; i < 2; i++)
      MSG("clock " << i << " " << total[i]);
  W = U;
  H = Z;
}

void run_master_slave_anls()
{
  if (MPI_RANK == 0) MSG("Master/Slave ANLS");

  PO_REQUIRE(vm, "data");
  PO_REQUIRE(vm, "k");
  PO_VARIABLE(data, vm, "data", std::string);
  PO_VARIABLE(k, vm, "k", int);

  using namespace arma;
  mat A;
  std::stringstream ss;
  ss << data << "_" << MPI_RANK;
  A.load(ss.str(), arma_binary);
  int m = A.n_rows, n = A.n_cols;
  
  srand(1+MPI_RANK);
  mat W_i(m, k), H_i(n, k);
  W_i.randu();
  H_i.randu();
  
  solve_master_slave_anls(A, W_i, H_i);
}
*/
